import os
import sys

import cv2
from PyQt5.QtCore import pyqtSlot, QSize
from PyQt5.QtGui import QImage, QPixmap, QIcon

from PyQt5.QtWidgets import QWidget, QLabel, QPushButton, QFileDialog, QApplication
from PIL import Image

from ultralytics import YOLO
import numpy as np
from PyQt5 import QtCore, QtWidgets

model = YOLO("D:\\Codes\\ultralytics-main\\yolov8n.pt")
model2 = YOLO("D:\\Codes\\ultralytics-main\\runs\\detect\\train36\\weights\\best.pt")


class App(QWidget):
    def __init__(self):
        super().__init__()

        self.img_path = None

        self.btn_photo = QPushButton('获取图片', self)
        self.btn_photo.setToolTip('点击后从电脑中读取图片')
        self.btn_photo.resize(200, 120)
        self.btn_photo.move(530, 50)
        self.btn_photo.clicked.connect(self.open_image)

        self.btn_photo_modify = QPushButton('处理图片', self)
        self.btn_photo_modify.setToolTip('点击后处理图片')
        self.btn_photo_modify.resize(200, 120)
        self.btn_photo_modify.move(530, 400)
        self.btn_photo_modify.clicked.connect(self.modify_image)

        self.label_h = 500
        self.label_w = 500
        self.label_img = QLabel(self)
        self.label_img.move(10, 50)
        self.label_img.setFixedSize(self.label_w, self.label_h)
        self.label_img.setText("show image")
        self.label_img.setStyleSheet("QLabel{background:white;}")
        self.label_img.setObjectName("image_show")

        self.label_img_new = QLabel(self)
        self.label_img_new.move(750, 50)
        self.label_img_new.setFixedSize(self.label_w, self.label_h)
        self.label_img_new.setText("show image")
        self.label_img_new.setStyleSheet("QLabel{background:white;}")
        self.label_img_new.setObjectName("image_show")

    @pyqtSlot()
    def open_image(self):
        the_image_url = QFileDialog.getOpenFileName()
        path = str(the_image_url).split("'")
        self.img_path = path[1]

        if len(the_image_url) == 0:
            return
        show = Image.open(self.img_path).convert("RGB")
        show = show.resize([self.label_w, self.label_h])

        showImage = QImage(np.array(show), np.shape(show)[1], np.shape(show)[0], QImage.Format_RGB888)
        self.label_img.setPixmap(QPixmap.fromImage(showImage))

    def modify_image(self):
        img = cv2.imread(self.img_path)
        path_new = "D:/new_image.jpg"

        result = model2.predict(source=img, epochs=100,
                               imgsz=640,
                               batch=1)[0]
        # 获取原始图像
        img = result.orig_img
        # 绘制预测结果
        annotated_image = result.plot()
        # 显示预测结果

        cv2.imwrite(path_new, annotated_image)

        show = Image.open(path_new).convert("RGB")
        show = show.resize([self.label_w, self.label_h])
        showImage = QImage(np.array(show), np.shape(show)[1], np.shape(show)[0], QImage.Format_RGB888)
        self.label_img_new.setPixmap(QPixmap.fromImage(showImage))


def main():
    app = QApplication(sys.argv)
    window = App()
    window.resize(1300, 700)
    window.move(300,300)
    window.setWindowTitle('图像处理')
    window.show()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()
